import chat_model
from langchain_core.messages import HumanMessage,trim_messages
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory,RunnableLambda

import pickle
import os

model = chat_model.get_model()

def token_counter(messages):
    total = 0
    for msg in messages:
        total += len(msg.content)
    return total

class ChatBotManager():
    memery_store = {}
    trimmer = trim_messages(
        max_tokens=256,  # 最大token数量
        strategy="last",  # 选择裁剪策略 "first"代表新加入的消息裁剪掉, "last"代表后加入的消息裁剪掉
        token_counter=token_counter,  # 计算token数量的方法
        include_system=False,  # 是否包含系统消息
        allow_partial=False,  # 是否允许一段message内携带的内容被裁剪掉
        start_on="human",  # 裁剪的起始位置
    )
    @classmethod
    def get_session_history(cls, session_id):
        if session_id not in cls.memery_store:
            cls.memery_store[session_id] = InMemoryChatMessageHistory()
        return cls.memery_store[session_id]

    @classmethod
    def save_memery(cls):
        with open("memery.pkl", "wb+") as f:
            pickle.dump(cls.memery_store, f)

    @classmethod
    def load_memery(cls):
        # 检查一下有没有这文件
        if os.path.exists("memery.pkl"):
            with open("memery.pkl", "rb") as f:
                cls.memery_store = pickle.load(f)
        else:
            cls.memery_store = {}

class ChatBot():
    def __init__(self, session_id: str):
        self.chat_chain = RunnableWithMessageHistory(
            ChatBotManager.trimmer|model,
            ChatBotManager.get_session_history,
        )
        self.config = {"configurable": {"session_id": session_id}}

    def chat(self, content):
        response = self.chat_chain.stream(
            [HumanMessage(content=content)],
            config=self.config,
        )
        for chunk in response:
            print(chunk.content, end="")

    def play(self):
        while True:
            text = input("请输入：")
            if text == "exit":
                break
            self.chat(text)
            print()

if __name__ == "__main__":
    session_id = "abc1"
    ChatBotManager.load_memery()
    bot = ChatBot(session_id)
    bot.play()
    ChatBotManager.save_memery()